import os, torch, random
import argparse
import json
from _code.Utils import createID
from gen_datadict import gen_datadict
from _code import getLogger

parser = argparse.ArgumentParser(description="entry point")
parser.add_argument('--ensemble-size', type=int, default=12,
                    help='size of ensemble')
parser.add_argument('--meta-class-size', type=int, default=12,
                    help='size of meta-classes')
parser.add_argument('--dataset', type=str, default='CAR', choices=['CAR', 'CUB', 'SOP'])
parser.add_argument('--backbone', type=str, default='resnet18', choices=['bninception', 'resnet18'])
parser.add_argument('--head-tail', action='store_true')
parser.add_argument('--attention', action='store_true')
parser.add_argument('--nb-epochs', type=int, default=12)
parser.add_argument('--batch-size', type=int, default=192)
args = parser.parse_args()

## train
# data_dict = torch.load('/pless_nfs/home/datasets/CAR/data_dict_emb.pth')
data_dict = gen_datadict(name=args.dataset)
dst = (f'results/{args.dataset}_{args.ensemble_size}ensemble_'
       f'{args.meta_class_size}metaclass_{args.nb_epochs}epochs').lower()
# if not os.path.exists(dst):
os.makedirs(dst)
class Config:
    pass
cfg = Config()
cfg.mode = 'train'
cfg.logf = str(os.path.join(dst, 'train_log.txt'))
cfg.debug = True
log = getLogger(cfg)

import datetime
start_time = datetime.datetime.now()
log.info(f"Start time: {start_time.strftime('%Y-%m-%d %H:%M:%S')}")

with open(os.path.join(dst, 'config.json'), 'w') as f:
    json.dump(vars(args), f)
# ID matrix
log.info('Creating ID')
ID = createID(args.meta_class_size, args.ensemble_size, len(data_dict['tra']))

from _code.Train import learn
x = learn(args, args.dataset, ID, dst, data_dict, num_epochs=args.nb_epochs, batch_size=args.batch_size)
x.run()
torch.save(ID, os.path.join(dst, 'ID.pth'))

end_time = datetime.datetime.now()
log.info(f"End time: {end_time.strftime('%Y-%m-%d %H:%M:%S')}")
total_tim = end_time - start_time  # timedelta
log.info(f'Total time: {total_tim}')
